# -*- coding: utf-8 -*-
# time: 2025/3/26 08:52
# file: llm_qwen_interl.py
# author: hanson

"""
conda activate pyenv
conda install llama-index langchain transformers sentence-transformers
llama-index-llms-langchain
"""

from llama_index.core import Settings, SimpleDirectoryReader, VectorStoreIndex
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.llms.langchain import LangChainLLM  # 使用LangChainLLM替代HuggingFaceLLM
from langchain.llms import HuggingFacePipeline  # 引入LangChain的HuggingFacePipeline

# 初始化一个HuggingFaceEmbedding对象，用于将文本转换为向量表示
embed_model = HuggingFaceEmbedding(
    model_name=r"E:\soft\embedding\Ceceliachenen\paraphrase-multilingual-MiniLM-L12-v2"
)
Settings.embed_model = embed_model

# 使用 HuggingFace 加载本地模型并通过LangChain包装
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline

tokenizer = AutoTokenizer.from_pretrained(r"E:\soft\model\qwen\Qwen\Qwen2___5-0___5B-Instruct")
model = AutoModelForCausalLM.from_pretrained(r"E:\soft\model\qwen\Qwen\Qwen2___5-0___5B-Instruct", trust_remote_code=True)
hf_pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer, trust_remote_code=True)

# 将HuggingFacePipeline包装为LangChainLLM
llm = LangChainLLM(llm=HuggingFacePipeline(pipeline=hf_pipeline))
Settings.llm = llm

# 从指定目录读取所有文档，并加载数据到内存中
documents = SimpleDirectoryReader(r"F:/workspace/py_project/intellect/llm/data", required_exts=[".md"]).load_data()

# 创建一个VectorStoreIndex，并使用之前加载的文档来构建索引
index = VectorStoreIndex.from_documents(documents)

# 创建一个查询引擎，这个引擎可以接收查询并返回相关文档的响应
query_engine = index.as_query_engine()
response = query_engine.query("MyBatis-Plus 支持数据库？")
print(response)
